import pandas as pd
import matplotlib as plt
import os


class BackTest():
    def __init__(self,r_path,cash=100000,):
        assert os.path.exists(r_path),'回测数据目录不存在'
        self.r_path = r_path
        self.cash = cash
        self.position = {}
        
    def get_data(self,ts_code,trade_date=None):
        '''
        1、ts_code: 股票代码
        2、trade_date: 交易日————如果trade_date为空，则返回ts_code的所有数据
        3、使用的是通信达导出的数据，一只股票数据为一个文件，文件名为XXXXXX.YY.csv，编码格式为gbk
        4、数据表格式：日期、开、高、低、收、量、额
        '''
        file_path = os.path.join(self.r_path,f'{ts_code}.csv')
        col_name = ['trade_date','open','high','low','close','vol','amount']
        df = pd.DataFrame()
        try:# 有些文件会有错误
            df = pd.read_csv(file_path,header=None,encoding='gbk',names=col_name)
        except pd.errors.ParserError:
            return df
        if len(df) == 0:
            return df
        if trade_date:
            return df[df['trade_date']==trade_date]
        # 去除异常数据、排好序、重新设置索引
        df.dropna(inplace=True)
        df.sort_values('trade_date',inplace=True)
        df.reset_index(drop=True)
        return df
    
    # 加仓购买
    def buy(self,ts_code,trade_date,num):
        '''股票代码、时间、手数'''
        df = self.get_data(ts_code,trade_date)
        if len(df) != 1:
            return False
        close_price = df.iloc[0]['close']
        if self.cash < close_price*100*num:
            return False
        self.cash -= close_price*100*num
        if ts_code in self.position.keys():
            self.position[ts_code] += num*100
        else:
            self.position[ts_code] = num*100
        return True
    
    # 卖出        
    def sell(self,ts_code,trade_date,num):
        '''股票代码、时间、手数'''
        df = self.get_data(ts_code,trade_date)
        if len(df) != 1:
            return False
        close_price = df.iloc[0]['close']
        if ts_code not in self.position:
            return False
        if self.position[ts_code] < 100*num:
            return False
        self.cash += close_price*100*num
        self.position[ts_code] -= num*100
        return True
    
    